//
// itertool - a polynormial and type-safe iterator framework in C++.
//
// Copyright (c) 2021-2025 Feng Cong
//
// Distributed under the MIT Software License
// See accompanying file LICENSE.txt or copy at
// https://opensource.org/licenses/MIT

#pragma once

#include <itertool/itertool_fwd.h>

#include <algorithm>
#include <iterator>
#include <numeric>
#include <vector>

namespace itertool {
namespace detail {

template <typename T>
class ChainIterable : public Iterable<T> {
 public:
  using ItemType = T;

  ChainIterable(std::initializer_list<Iterator<T>> list) {
    Init(list.begin(), list.end(), list.size());
  }

  explicit ChainIterable(Iterator<Iterator<T>> iterators) {
    assert(iterators.IsFinite());
    Init(iterators.begin(), iterators.end(), iterators.LengthHint());
  }

  template <typename InputIter>
  ChainIterable(InputIter first, InputIter last) {
    Init(first, last, std::distance(first, last));
  }

  bool HasNext() override { return now_ < data_.size(); }

  ItemType GetNow() override {
    assert(HasNext());
    return data_[now_].GetNow();
  }

  void Advance() override {
    assert(HasNext());
    data_[now_].Advance();
    if (!data_[now_].HasNext()) {
      ++now_;
    }
  }

  bool IsResetable() override {
    return std::all_of(data_.begin(), data_.end(), &Tr::IsResetable);
  }

  bool IsFinite() override {
    return std::all_of(data_.begin(), data_.end(), &Tr::IsFinite);
  }

  void Reset() override {
    now_ = 0;
    std::for_each(data_.begin(), data_.end(), &Tr::Reset);
  }

  size_t LengthHint() override {
    return std::accumulate(
        data_.begin() + now_, data_.end(), size_t(0),
        [](size_t sum, Iterator<T> it) { return sum + it.LengthHint(); });
  }

 private:
  using Tr = IterableTraits<Iterator<T>>;

  template <typename InputIter>
  void Init(InputIter first, InputIter last, size_t length_hint) {
    data_.reserve(length_hint);
    std::copy_if(first, last, std::back_inserter(data_), &Tr::HasNext);
  }

  size_t now_ = 0;
  IteratorVector<T> data_;
};

template <typename T>
class InfiniteChainIterable : public Iterable<T> {
 public:
  using ItemType = T;

  explicit InfiniteChainIterable(Iterator<Iterator<T>> data) : data_(data) {
    FindNextNonEmpty();
  }

  bool HasNext() override { return !now_.IsEmpty(); }

  void Advance() override {
    assert(HasNext());
    now_.Advance();
    if (!now_.HasNext()) {
      FindNextNonEmpty();
    }
  }

  ItemType GetNow() override {
    assert(HasNext());
    return now_.GetNow();
  }

 private:
  void FindNextNonEmpty() {
    while (data_.HasNext()) {
      Iterator<T> val = data_.GetNow();
      data_.Advance();
      if (val.HasNext()) {
        now_ = val;
        return;
      }
    }
    now_.Clear();
  }

  Iterator<T> now_;
  Iterator<Iterator<T>> data_;
};

}  // namespace detail

template <typename T>
Iterator<T> Chain(std::initializer_list<Iterator<T>> list) {
  return MakeIterator<detail::ChainIterable<T>>(list);
}

template <typename T>
Iterator<T> Chain(Iterator<Iterator<T>> iterators) {
  return iterators.IsFinite()
             ? MakeIterator<detail::ChainIterable<T>>(iterators)
             : MakeIterator<detail::InfiniteChainIterable<T>>(iterators);
}

template <
    typename InputIter,
    typename _ValueType = typename std::iterator_traits<InputIter>::value_type,
    typename _ItemType = typename _ValueType::ItemType>
Iterator<_ItemType> Chain(InputIter first, InputIter last) {
  return MakeIterator<detail::ChainIterable<_ItemType>>(first, last);
}

}  // namespace itertool
